# -*- coding:utf-8 -*-
import glob # 导入操作文件的模块
import xml.etree.ElementTree as ET # 导入解析XML文件格式的模块
from tqdm import tqdm # 导入进度条显示模块
import numpy as np # 导入矩阵计算模块
import matplotlib.pyplot as plt # 导入绘图模块
from sklearn.cluster import KMeans # 使用sklearn做kmeans聚类分析

def iou(box, clusters):
    """
    Usage：
    Calculates the Intersection over Union (IoU) between a box and k clusters.
    box: tuple or array, shifted to the origin (i. e. width and height)
    clusters: numpy array of shape (k, 2) where k is the number of clusters
    numpy array of shape (k, 0) where k is the number of clusters
    """
    x = np.minimum(clusters[:, 0], box[0])
    y = np.minimum(clusters[:, 1], box[1])
    if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
        raise ValueError("Box has no area")
    intersection = x * y
    box_area = box[0] * box[1]
    cluster_area = clusters[:, 0] * clusters[:, 1]
    iou_ = intersection / (box_area + cluster_area - intersection)
    return iou_

def avg_iou(boxes, clusters):
    """
    Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param clusters: numpy array of shape (k, 2) where k is the number of clusters
    :return: average IoU as a single float
    """
    return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])

def translate_boxes(boxes):
    """
    Translates all the boxes to the origin.
    :param boxes: numpy array of shape (r, 4)
    :return: numpy array of shape (r, 2)
    """
    new_boxes = boxes.copy()
    for row in range(new_boxes.shape[0]):
        new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])
        new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])
    return np.delete(new_boxes, [0, 1], axis=1)

def kmeans(boxes, k=9, dist=np.median):
    """
    Calculates k-means clustering with the Intersection over Union (IoU) metric.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param k: number of clusters
    :param dist: distance function
    :return: numpy array of shape (k, 2)
    """
    rows = boxes.shape[0]
    distances = np.empty((rows, k))
    last_clusters = np.zeros((rows,))
    np.random.seed()
    # the Forgy method will fail if the whole array contains the same rows
    clusters = boxes[np.random.choice(rows, k, replace=False)]
    while True:
        for row in range(rows):
            distances[row] = 1 - iou(boxes[row], clusters)
        nearest_clusters = np.argmin(distances, axis=1)
        if (last_clusters == nearest_clusters).all():
            break
        for cluster in range(k):
            clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)
        last_clusters = nearest_clusters
    return clusters

def show_cluster(data,cluster, max_points=2000):
    if len(data) > max_points:
        idx = np.random.choice(len(data),max_points)
        data = data[idx]
    plt.scatter(data[:, 0], data[:, 1], s=5, c='lavender')
    plt.scatter(cluster[:, 0], cluster[:, 1], c='red', s=100, marker="^")
    plt.xlabel("Width")
    plt.ylabel("Height")
    plt.title("Bounding and anchor distribution")
    plt.savefig("cluster.png")
    plt.show()

def show_w_h(data,cluster,bins=50):
    if data.dtype != np.float32:
        data = data.astype(np.float32)
    width = data[:, 0]
    height = data[:, 1]
    ratio = height / width

    plt.figure(1, figsize=(20, 6))
    plt.subplot(131)
    plt.hist(width, bins=bins, color='blue', rwidth=0.8)
    plt.xlabel('Width')
    plt.ylabel('Number')
    plt.grid(True, linestyle='-.')
    plt.title('Distribution of Width')

    plt.subplot(132)
    plt.hist(height, bins=bins, color='green', rwidth=0.8)
    plt.xlabel('Height')
    plt.ylabel('Number')
    plt.grid(True, linestyle='-.')
    plt.title('Distribution of Height')

    plt.subplot(133)
    plt.hist(ratio, bins=bins, color='magenta', rwidth=0.8)
    plt.xlabel('Height / Width')
    plt.ylabel('Number')
    plt.grid(True, linestyle='-.')
    plt.title('Distribution of aspect ratio[Height / Width]')
    plt.savefig("shape-distribution.png")
    plt.show()

def sort_cluster(cluster):
    if cluster.dtype != np.float32:
        cluster = cluster.astype(np.float32)
    area = cluster[:, 0] * cluster[:, 1]
    ratio = cluster[:, 1:2] / cluster[:, 0:1]
    return np.concatenate([cluster, ratio], axis=-1)

def load_dataset(path):
    dataset = []
    for xml_file in tqdm(glob.glob("{}/*xml".format(path))):
        # print(xml_file)
        tree = ET.parse(xml_file)

        height = int(tree.findtext("./size/height"))
        width = int(tree.findtext("./size/width"))

        for obj in tree.iter("object"):
            xmin = int(float(obj.findtext("bndbox/xmin"))) / width
            ymin = int(float(obj.findtext("bndbox/ymin"))) / height
            xmax = int(float(obj.findtext("bndbox/xmax"))) / width
            ymax = int(float(obj.findtext("bndbox/ymax"))) / height

            dataset.append([xmax - xmin, ymax - ymin])

    return np.array(dataset)

if __name__ == "__main__":
    ANNOTATIONS_PATH = r"E:\\Bullet\\VOCdevkit\\VOC2007\\Annotations"
    #以正斜杠/这种形式可以防止反斜杠带来的转义错误
    CLUSTERS = 9
    BBOX_NORMALIZE = True
    data = load_dataset(ANNOTATIONS_PATH)
    #print(ANNOTATIONS_PATH)
    out = kmeans(data, k=CLUSTERS)
    print(out.shape)
    out_sorted = sort_cluster(out)
    #print(out_sorted)
    # if out.dtype != np.float32:
    #     cluster = out.astype(np.float32)
    show_cluster(data, out, max_points=5000)
    show_w_h(data, out, bins=50)